#!/usr/bin/python

# -*- coding: utf-8 -*-

from gevent import monkey
monkey.patch_all()
from gevent.pool import Pool
from gevent import socket
import BaseHTTPServer
#from gevent import socket
from gevent.server import StreamServer
from Crypto.Cipher import ARC4
import urlparse
from gevent.select import select
#import zlib
import httplib
import traceback
import ConfigParser

import util

proxy_host = "127.0.0.1"
proxy_port =  8038


headers_to_server={
    "User-Agent":"Mozilla/5.0 (Windows NT 6.1) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/34.0.1847.137 Safari/537.36",
    "Accept":"*/*",
    #"Transfer-Encoding" : "deflate",
    "Accept-Language":"en-US,en",
    "Accept-Encoding":"deflate,gzip",
    "Host":"djvps-1.softether.net",
    "Connection":"keep-alive",
    #"Content-Encoding":"deflate",
    "Content-Type" : "application/json",
#content_encoding = "chunked"
}


class Handler(BaseHTTPServer.BaseHTTPRequestHandler):
    server_version = "pws"
    protocol_version = "HTTP/1.1"
    close_connection  = 0
    def __init__(self, request, addr):
        self.request = request
        self.client_address = addr
        #self.rfile = request.makefile("rb", -1)
        #self.wfile = request.makefile("wb", 0)
        self.rsock = None
        self.conn = None
        self.setup()

        self.handle()
        self.finish()
        
    def do_GET(self):
        if self.path.startswith("/"):
            msg = "hello, world"
            self.send_response(200)
            self.send_header("Content-Length", len(msg))
            self.log_request(200, len(msg))
            self.send_header("Connection", "close")
            self.end_headers()
            self.wfile.write(msg)
            return

        self.do_METHOD()

    def do_METHOD(self):
        self.lkey,self.sessionid = util.random_key(password)

        enc = ARC4.new(self.lkey)
        
        body_len = int(self.headers.get("content-length",0))
        body = ''

        # read request body
        if body_len:
            body = self.rfile.read(body_len)

        data_to_send = "%s %s %s\r\n" % (self.command, self.path, self.request_version)

        for k,v in self.headers.items():
            if k.lower() != "proxy-connection":
                data_to_send += "%s: %s\r\n" % (k, v)
        data_to_send += 'Content-Length: %d\r\n' % len(body)
        data_to_send += "\r\n"
        data_to_send += body

        enc_data = enc.encrypt(data_to_send)

        headers = headers_to_server.copy()
        headers["Cookie"] = "sessionid=%s" % self.sessionid
        headers["Content-Length"] = len(enc_data)
        #self.conn = None
        try:
            if not self.conn:
                self.conn = httplib.HTTPConnection(proxy_host, proxy_port)
            self.conn.request("POST",util.random_url(), enc_data, headers)
            response = self.conn.getresponse()
        except:
            traceback.print_exc()
            self.send_response(500)
            self.send_header("Proxy-Connection", "close")
            self.end_headers()
            self.log_request(500)
            if self.conn:
                self.conn.close()
            return

        # may be server error
        if response.status != 200:
            self.send_response(response.status)
            self.send_header("Proxy-Connection","close")
            self.end_headers()
            self.log_request(response.status)
            self.conn.close()
            return

        # get key
        sid = response.getheader("Set-Cookie").split("=", 1)[1]
        rkey = util.decode_key(sid, password)
        
        # body
        data = response.read()

        # decrypt
        enc = ARC4.new(rkey)
        d = enc.encrypt(data)

        #self.send_response(response.status)
        self.wfile.write(d)

        self.log_request(200, len(d))
        #self.conn.close()
        self.close_connection = 0
        
    def do_POST(self):
        if self.path.startswith("/"):
            self.send_response(400)
            self.send_header("Connection","close")
            self.end_header()
            return
        self.do_METHOD()

    def do_CONNECT(self):

        key,sessionid = util.random_key(password)

        enc = ARC4.new(key)

        data_to_send = "%s %s %s\r\n" % (self.command, self.path, self.request_version)
        #print data_to_send
        for k,v in self.headers.items():
            data_to_send += '%s: %s\r\n' % (k, v)
        data_to_send += "\r\n"

        enc_data = enc.encrypt(data_to_send)
        if not self.rsock:
            self.rsock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            self.rsock.connect((proxy_host, proxy_port))
        #print "connect to proxy success"

        #rfile = self.rsock.makefile("rb", 0)
        #wfile = self.rsock.makefile("wb", 0)
        try:
            self.rsock.send("POST %s HTTP/1.1\r\n" % util.random_url())
            for k,v in headers_to_server.items():
                self.rsock.send("%s: %s\r\n" % (k.title(), v))
            #self.rsock.send("Content-Length: %d\r\n" % len(enc_data))
            self.rsock.send("Transfer-Encoding: chunked\r\n")
            self.rsock.send("Content-Encoding: deflate\r\n")
            self.rsock.send("Cookie: sessionid=%s\r\n" % sessionid)
            self.rsock.send("\r\n")
            self.rsock.send("%x\r\n%s\r\n" % (len(enc_data), enc_data))
        except:
            traceback.print_exc()
            self.rsock.close()
            self.rsock = None
            return
        #wfile.flush()

        rfile = self.rsock.makefile("rb", -1)
        firstline = rfile.readline()
        if not firstline:
            self.rsock.close()
            self.rsock = None
            return
        version, status, reason = firstline.strip(" \r\n").split(" ", 2)

        self.log_message("%s %s %s -", self.command, self.path, status) 


        remote_key = ""
        while True:
            line = rfile.readline()
            if not line:
                break
            if line == '\r\n' or line == '\n':
                break
            if "Set-Cookie" in line:
                raw_key = line.strip(" \r\n").split(":")[1].strip(" \r\n").split("=",1)[1]
                remote_key = util.decode_key(raw_key, password)

        if int(status) != 200:
            self.send_response(int(status))
            self.send_header("Proxy-Connection", "close")
            self.end_headers()
            rfile.close()
            
            return
        
        #rfile.close()
        if not remote_key:
            raise TypeError("get key failed")
            
        r_enc = ARC4.new(remote_key)
        #print "begin to forward"
        while True:
            r,_,_ = select([self.request, self.rsock], [], [])
            if self.request in r:
                data = self.request.recv(1024)
                if not data:
                    break
                enc_data = enc.encrypt(data)    
                self.rsock.send("%x\r\n%s\r\n" % (len(enc_data), enc_data))
            if self.rsock in r:
                l = rfile.readline().strip(" \r\n")
                if not l:
                    break
                l = int(l, 16)
                data = ''
                unfinished = l
                while unfinished > 0:
                    d = rfile.read(unfinished)
                    unfinished -= len(d)
                    data += d
                rfile.readline()    
                self.request.send(r_enc.encrypt(data))
        #wfile.close()
        rfile.close()
        #self.rsock.close()
    
    def address_string(self):
        return "%s:%d" % self.client_address

    def send_response(self, code, message=None):
        if message is None:
            if code in self.responses:
                message = self.responses[code][0]
            else:
                message = ''
        self.wfile.write("%s %d %s\r\n" % (
            self.protocol_version, code, message))
        self.send_header("Server", self.server_version)
        self.send_header("Date", self.date_time_string())

    def send_header(self, keyword, value, fp=None):
        out = self.wfile
        if fp:
            out = fp
        out.write("%s: %s\r\n" % (keyword, value))
        if keyword.lower() == 'connection' or keyword.lower() == 'proxy-connection':
            if value.lower() == 'close':
                self.close_connection = 1
            elif value.lower() == 'keep-alive':
                self.close_connection = 0

    def end_headers(self, fp=None):
        out = self.wfile
        if fp:
            out = fp
        out.write("\r\n")


if __name__ == "__main__":
    cfg=ConfigParser.ConfigParser()
    cfg.read("local.ini")
    pool = Pool(100)
    host = cfg.get("local","listen")
    port = cfg.getint("local","port")
    proxy_host = cfg.get("local","remote_host")
    proxy_port = cfg.getint("local","remote_port")
    password = cfg.get("local","password")
    srv = StreamServer((host, port), Handler,spawn=pool)
    srv.serve_forever()

